/*
 * Copyright 2022 The Open Islands Authors. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.openislands.oi.master.scheduler;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.google.common.base.CaseFormat;
import com.google.common.base.Strings;
import org.openislands.oi.config.MyNodeInfo;
import org.openislands.oi.constant.*;
import org.openislands.oi.manager.CommonJobManager;
import org.openislands.oi.scheduling.DAG;
import org.openislands.oi.scheduling.DAGException;
import org.openislands.oi.scheduling.DAGOperator;
import org.openislands.oi.scheduling.StateEngine;
import org.openislands.oi.error.LogicException;
import org.openislands.oi.pojo.dto.JobDTO;
import org.openislands.oi.pojo.dto.TaskDTO;
import org.openislands.oi.util.CustomerLogger;
import org.openislands.oi.util.CustomerLoggerFactory;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.lang.reflect.Method;
import java.util.*;
import java.util.stream.Collectors;

@Component
public class DAGScheduler implements InitializingBean {
    private static final CustomerLogger schedulingLogger = CustomerLoggerFactory.getLogger(LogTraceType.JOB, LogType.SCHEDULING);

    @Resource
    private CommonJobManager commonJobManager;
    @Resource(name = "MasterTaskScheduler", type = TaskScheduler.class)
    private TaskScheduler taskScheduler;
    @Resource
    private MyNodeInfo myNodeInfo;
    @Resource
    private StateEngine stateEngine;
    @Resource
    private JobIdGenerator jobIdGenerator;

    private final Map<JobState, Method> stateHandleEvent;

    @Override
    public void afterPropertiesSet() throws Exception {
        this.init();
    }

    public DAGScheduler() throws NoSuchMethodException {
        this.stateHandleEvent = new HashMap<>();
    }

    private void init() throws NoSuchMethodException {
        schedulingLogger.info("start to init DAGScheduler");
        for (JobState state: StateEngine.getAllJobStates()){
            String methodName = String.format("handle%sJob", CaseFormat.UPPER_UNDERSCORE.to(CaseFormat.UPPER_CAMEL, state.name()));
            try {
                Method method = this.getClass().getDeclaredMethod(methodName, JobDTO.class);
                this.stateHandleEvent.put(state, method);
            }catch (Exception e){
                schedulingLogger.error("not found {} state handle method {}", state, methodName);
                throw e;
            }
        }
        schedulingLogger.info("init DAGScheduler done");
    }

    public JobDTO submitJob(JobDTO jobDTO){
        if (Strings.isNullOrEmpty(jobDTO.getJobId())){
            jobDTO.setJobId(jobIdGenerator.nextId());
        }
        //todo: using parser to get runtimeConf by conf and confOverwrite
        jobDTO.setRuntimeConf(jobDTO.getConf());
        DAG dag = this.getDAG(jobDTO);
        jobDTO.setState(JobState.READY);
        jobDTO.setNextState(JobState.READY);
        boolean st = taskScheduler.createJob(dag, jobDTO);
        if (st){
            this.createTasks(jobDTO, dag);
            return jobDTO;
        }else {
            throw new LogicException(MessageCodeEnum.DAG_SCHEDULING_ERROR, "create job error");
        }
    }

    private void fillResourceParams(JobDTO jobDTO){
        //todo: engine address
    }

    private void saveConfs(JobDTO jobDTO){
        //todo: save conf/confOverwrite/runtimeConf
    }

    public Boolean stopJob(String jobId){
        JobDTO jobDTO = commonJobManager.queryJob(jobId);
        if (jobDTO == null){
            throw new LogicException(MessageCodeEnum.DATA_NOT_FOUND_ERROR, "can not found this job");
        }
        return taskScheduler.stopJob(jobDTO, JobState.CANCELED);
    }

    public void schedulingJobs(){
        for (JobState state : stateEngine.getUnfinishedJobStates()) {
            schedulingLogger.info("start to schedule {} state job", state);
            try{
                List<JobDTO> handleJobs = commonJobManager.queryJobByState(state, myNodeInfo.getNodeId());
                schedulingLogger.info("get {} jobs", handleJobs.size());
                for (JobDTO jobDTO:handleJobs){
                    try{
                        schedulingLogger.info(jobDTO.getJobId(), "handle job");
                        this.handleJob(jobDTO);
                        schedulingLogger.info(jobDTO.getJobId(), "handle job done");
                    }catch (Exception e){
                        schedulingLogger.error(jobDTO.getJobId(), "handle job error", e);
                    }
                }
            }catch (Exception e){
                schedulingLogger.error("schedule {} state job error", state, e);
            }
            schedulingLogger.info("schedule {} state job done", state);
        }
    }
    public void handleJob(String jobId){
        JobDTO jobDTO = commonJobManager.queryJob(jobId);
        if(jobDTO == null){
            throw new LogicException(MessageCodeEnum.DATA_NOT_FOUND_ERROR, "can not found job");
        }
        this.handleJob(jobDTO);
    }

    private void handleJob(JobDTO jobDTO){
        try{
            this.stateHandleEvent.get(jobDTO.getState()).invoke(this, jobDTO);
        }catch (Exception e){
            throw new LogicException(e, MessageCodeEnum.SYSTEM_ERROR, String.format("handle job %s with %s state error", jobDTO.getJobId(), jobDTO.getState()));
        }
    }
    public void handleReadyJob(JobDTO jobDTO){
        DAG dag = this.getDAG(jobDTO);
        this.createTasks(jobDTO, dag);
    }

    public void createTasks(JobDTO jobDTO, DAG dag){
        schedulingLogger.info(jobDTO.getJobId(), "start to create tasks");
        boolean allSuccess = true;
        for (DAGOperator dagOperator: dag.getOperator()){
            if (!taskScheduler.createTask(jobDTO, dag, dagOperator)){
                allSuccess = false;
                schedulingLogger.error(jobDTO.getJobId(), "create task: {} failed", dagOperator.getName());
                break;
            }
        }
        JobState toState;
        if (allSuccess){
            schedulingLogger.info(jobDTO.getJobId(), "create tasks success");
            toState = JobState.SUBMITTED;
        }else {
            schedulingLogger.info(jobDTO.getJobId(), "create tasks not success, submission failed");
            toState = JobState.SUBMISSION_FAILED;
        }
        if(!taskScheduler.updateJobState(jobDTO, toState)){
            throw new LogicException(MessageCodeEnum.DAG_SCHEDULING_ERROR, String.format("sync %s state failed", JobState.SUBMITTED));
        }
    }

    public void handleSubmittedJob(JobDTO jobDTO){
        DAG dag = DAG.newDAG(jobDTO.getRuntimeConf());
        List<DAGOperator> roots = dag.filterOperatorRoot();
        DAGOperator root = roots.get(0);
        JobState newJobState;
        if(taskScheduler.submitTask(jobDTO, root.getName(), 0)){
            newJobState = JobState.RUNNING;
        }else {
            newJobState = JobState.SUBMISSION_FAILED;
            schedulingLogger.error(jobDTO.getJobId(), "submit job root task {} failed", root.getName());
        }
        boolean st = taskScheduler.updateJobState(jobDTO, newJobState);
        if (!st){
            schedulingLogger.error(jobDTO.getJobId(), "update job state to {} failed", newJobState);
        }
    }

    public void handleSubmissionFailedJob(JobDTO jobDTO) {
        throw new LogicException(MessageCodeEnum.DAG_SCHEDULING_ERROR, "no processing method");
    }

    public void handleUpForRetryJob(JobDTO jobDTO) {
        throw new LogicException(MessageCodeEnum.DAG_SCHEDULING_ERROR, "no processing method");
    }

    public void handleRunningJob(JobDTO jobDTO){
        TaskScheduler.TaskSchedulingResult taskSchedulingResult = taskScheduler.schedulingTasks(jobDTO.getJobId());
        JobState newJobState = this.calculateJobState(taskSchedulingResult.getStates(), DAGSchedulingState.HAVE_NEXT);
        String taskStatesMsg = taskSchedulingResult.getTaskStates().entrySet().stream().map(e -> String.format("%s:%s", e.getKey(), e.getValue().name())).collect(Collectors.joining(","));
        int progress = this.calculateJobProgress(taskSchedulingResult.getStates());
        schedulingLogger.info(jobDTO.getJobId(), "job current {}, calculate job {} new state by task states: {}", jobDTO.getState(), newJobState, taskStatesMsg);
        if (progress - jobDTO.getProgress() > 0){
            jobDTO.setProgress(progress);
            taskScheduler.updateJob(jobDTO);
        }
        if (jobDTO.getState() != newJobState){
            taskScheduler.updateJobState(jobDTO, newJobState);
            jobDTO.setState(newJobState);
            if (stateEngine.isFinishedJobState(newJobState)){
                this.handleJob(jobDTO);
            }
        }
        if (!stateEngine.isFinishedJobState(jobDTO.getState()) && !taskSchedulingResult.getReadyTasks().isEmpty()){
            DAG dag = this.getDAG(jobDTO);
            Map<String, List<String>> upstream = dag.getUpstream();
            for (TaskDTO readyTask: taskSchedulingResult.getReadyTasks()){
                boolean upstreamSuccess = true;
                for (String taskName: upstream.get(readyTask.getTaskName())){
                    if (taskSchedulingResult.getTaskStates().get(taskName) != TaskState.SUCCESS){
                        upstreamSuccess = false;
                    }
                }
                if (upstreamSuccess){
                    if(taskScheduler.submitTask(jobDTO, readyTask.getTaskName(), readyTask.getTaskVersion())){
                        schedulingLogger.info(jobDTO.getJobId(), "submit {} {} task", readyTask.getTaskName(), readyTask.getTaskVersion());
                    }
                }
            }
        }
    }

    public void handleCanceledJob(JobDTO jobDTO) {
        this.finishJob(jobDTO);
    }

    public void handleTimeoutJob(JobDTO jobDTO) {
        this.finishJob(jobDTO);
    }

    public void handleRunFailedJob(JobDTO jobDTO) {
        this.finishJob(jobDTO);
    }

    public void handleSuccessJob(JobDTO jobDTO) {
        //todo: save model
        this.finishJob(jobDTO);
    }

    private void finishJob(JobDTO jobDTO){
        schedulingLogger.info(jobDTO.getJobId(), "start to finish job");
        this.stopJob(jobDTO.getJobId());
        this.cleanJob(jobDTO.getJobId());
        schedulingLogger.info(jobDTO.getJobId(), "finish job with {} state done", jobDTO.getState());
        //schedulingLogger.info("job {} finish with {} state", jobDTO.getJobId(), jobDTO.getState());
    }

    public Boolean cleanJob(String jobId){
        return true;
    }

    private DAG getDAG(JobDTO jobDTO){
        DAG dag = DAG.newDAG(jobDTO.getRuntimeConf());
        try {
            dag.check();
            dag.calDependency();
        } catch (DAGException | JsonProcessingException e) {
            throw new RuntimeException(e);
        }
        return dag;
    }

    private JobState calculateJobState(Set<TaskState> taskStates, DAGSchedulingState dagSchedulingState){
        if (taskStates.size() == 1){
            return stateEngine.getTaskStateMapJob(taskStates.iterator().next());
        }
        //have running
        if (taskStates.contains(TaskState.RUNNING)){
            // rule 2
            return JobState.RUNNING;
        }
        // have interrupt
        List<TaskState> withInterrupt = stateEngine.intersectionWithTaskStates(taskStates, stateEngine.getInterruptTaskStates());
        if (!withInterrupt.isEmpty()){
            return stateEngine.getTaskStateMapJob(withInterrupt.get(0));
        }
        // success/not interrupt
        if (taskStates.contains(TaskState.SUCCESS)){
            return JobState.RUNNING;
        }
        // different unfinished
        List<TaskState> withUnfinished = stateEngine.intersectionWithTaskStates(taskStates, stateEngine.getUnfinishedTaskStates());
        if (!withUnfinished.isEmpty()){
            return stateEngine.getTaskStateMapJob(withUnfinished.get(0));
        }
        String statesDesc = taskStates.stream().map(Enum::name).collect(Collectors.joining(","));
        throw new LogicException(MessageCodeEnum.DAG_SCHEDULING_ERROR, String.format("no rule to calculating job state, task states: %s", statesDesc));
    }

    private Integer calculateJobProgress(Set<TaskState> taskStates){
        int finished = 0;
        for (TaskState state: taskStates){
            if (stateEngine.isFinishedTaskState(state)){
                finished += 1;
            }
        }
        return finished * 100 / taskStates.size();
    }
}
